[Vertex AI] Add countTokens support for Developer API via VinF
#14644
Conversation
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback. |
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for countTokens using the Developer API through the Vertex AI in Firebase backend. The changes include modifications to the GenerativeModel, VertexAI, and CountTokensRequest classes to handle the different model resource name formats required by the various backends. Integration tests have also been updated to ensure no regressions. Overall, the changes seem well-structured and address the specific edge case described in the pull request description.
Summary of Findings
- Model Name Handling: The introduction of the
modelNameproperty inGenerativeModeland its usage in constructing thegenerateContentRequestModelResourceNameseems like a good approach to handle the different model name formats. However, it's crucial to ensure that this new property is always correctly initialized and used consistently throughout the codebase. - Code Duplication: The
developerModelResourceNamefunction inVertexAI.swifthas similar logic for.firebaseVertexAIStagingand.firebaseVertexAIProd. Consider refactoring to reduce code duplication. - Test Coverage: The integration tests have been updated, but it's important to ensure that all edge cases and different configurations are thoroughly tested to prevent regressions.
- Documentation: The documentation for the
modelResourceNameparameter in theGenerativeModelinitializer is comprehensive and helpful. Ensure that all public APIs are well-documented.
Merge Readiness
The pull request introduces important functionality for handling a specific edge case with the countTokens API. The code appears to be well-structured and addresses the issue effectively. However, before merging, it's recommended to address the potential code duplication in developerModelResourceName and ensure comprehensive test coverage for all configurations. I am unable to directly approve this pull request, and other reviewers should review and approve this code before merging.
FirebaseVertexAI/Tests/TestApp/Tests/Integration/CountTokensIntegrationTests.swift
Show resolved
Hide resolved
Added support for
countTokensusing the Developer API through the Vertex AI in Firebase (firebasevertexai.googleapis.com) backend. This combination is an edge case where the model resource name (e.g.,projects/my-project/models/gemini-2.0-flash) in thecountTokensrequest URL differs from the resource name in the request fieldCountTokensRequest.generateContentRequest.name. The integration tests have been updated to perform acountTokensrequest using bothfirebasevertexai.googleapis.com(Vertex AI and Developer API) andgenerativelanguage.googleapis.com(Developer API only) to ensure no regressions.#no-changelog